#!/usr/bin/python2
# -*- coding: utf-8 -*-

# Crypto
#AES has a fixed block size of 128 bits,
#so "AES 256" means 128 bit blocks, 256(MAX) bit key, 14 rounds.
from Crypto.Cipher import AES, Blowfish
#aes-256-cfb -- key 32 bytes (SHA256)

# TODO: 
# USE M2Crypto against Crypto
# Crypto: t_enc: 29.7992401123
# M2Crypto: t_enc : 4,861

import hashlib
from hashlib import sha256 as SHA256
from hashlib import sha512 as SHA512
from hashlib import md5 as MD5
from binascii import hexlify, unhexlify
import sys
import os

# M2Crypto
from M2Crypto import EVP

# (512, 1024, 4*1024, 16*1024, 64*1024)
CHUNKSIZE = ('512', '1024', '4096', '16384', '65536')
MODE = {'CFB': AES.MODE_CFB, 'PGP': AES.MODE_PGP}
CIPHER = {'AES': [AES, 16], 'BLOWFISH': [Blowfish, 8]}
KEYSIZE = {'128': MD5, '256': SHA256}


def pass2key(password, message='', keysize=256):
    keysize = str(keysize)
    assert KEYSIZE.has_key(keysize), 'Bad keysize. Use 128 or 256'
    message = str(message)
    return HMAC.new(str(password),
                        msg=message,
                        digestmod=KEYSIZE.get(keysize)).digest()


'''
EXAMPLE:
        # the crypto algorithms are unicode unfriendly
        if isinstance(password, unicode):
            password = password.encode('utf8')
        
        # get 256 bit random encryption salt
        self.salt = os.urandom(32)
        # derive 256 bit encryption key using the pbkdf2 standard
        key = EVP.pbkdf2(password, self.salt, iter=self.iterations, keylen=32)
        
        # Derive encryption key and HMAC key from it
        # See Practical Cryptography section 8.4.1.
        hmacKey = SHA256(key + 'MAC').digest() # or MD5(key + 'MAC').digest()
        # encKey используется для шифрования данных (сессионный ключ)
        encKey = SHA256(key + 'SessCrypt').digest() # or MD5(key + 'SessCrypt')
        del key


        # mac хранится для проверки правильности введеного login-password
        mac = HMAC.new(hmacKey,
                       msg=self.salt,
                       digestmod=SHA512).digest()
'''


def pkcs5_pad(data, blklen=8):
    pad=(8-(len(data)%8))
    return data+chr(pad)*pad

def pkcs7_pad(data, blklen):
    if blklen>255:
        raise ValueError, 'illegal block size'
    pad=(blklen-(len(data)%blklen))
    return data+chr(pad)*pad
    
class HMAC(object):
    """RFC2104 HMAC class.

    This supports the API for Cryptographic Hash Functions (PEP 247).
    """

    def __init__(self, key, msg=None, digestmod=None):
        """Create a new HMAC object.

        key:       key for the keyed hash object.
        msg:       Initial input for the hash, if provided.
        digestmod: A module supporting PEP 247. Defaults to the md5 module.
        """

        if digestmod is None:
            digestmod = 'MD5'

        #digest = hashlib.new(digestmod)
        self.outer = hashlib.new(digestmod)
        self.inner = hashlib.new(digestmod)
        try:
            self.digest_size = digestmod.digest_size
        except AttributeError:
            self.digest_size = len(self.outer.digest())

        try:
            # The block size is 128 bytes for SHA384 and SHA512 and 64 bytes
            # for the others hash function
            blocksize = digestmod.block_size
        except AttributeError:
            blocksize = 64

        ipad = 0x36
        opad = 0x5C

        if len(key) > blocksize:
            key = digestmod.new(key).digest()

        key = key + chr(0) * (blocksize - len(key))
        self.outer.update(''.join([chr(opad ^ ord(x))
                                    for x in key]))
        self.inner.update(''.join([chr(ipad ^ ord(x))
                                    for x in key]))
        if (msg):
            self.update(msg)

##    def clear(self):
##        raise NotImplementedError, "clear() method not available in HMAC."

    def update(self, msg):
        """Update this hashing object with the string msg.
        """
        self.inner.update(msg)

    def copy(self):
        """Return a separate copy of this hashing object.

        An update to this copy won't affect the original object.
        """
        other = HMAC(b(""))
        other.digestmod = self.digestmod
        other.inner = self.inner.copy()
        other.outer = self.outer.copy()
        return other

    def digest(self):
        """Return the hash value of this hashing object.

        This returns a string containing 8-bit data.  The object is
        not altered in any way by this function; you can continue
        updating the object after calling this function.
        """
        h = self.outer.copy()
        h.update(self.inner.digest())
        return h.digest()

    def hexdigest(self):
        """Like digest(), but returns a string of hexadecimal digits instead.
        """
        return hexlify(self.digest())


class Crypts(object):
    def __init__(self, password='', login='', key='',
                 cipher='AES', ciphermode="CFB",
                 keysize='256', blocksize='1024'):
        """ Encrypts/Decrypts a file using AES (CBC mode) with the
        given key.
        password:
            The secret string.
        login:
            Adding message to the password for more secure
        or
        key:
            String keysize long
        keysize:
            Size of key. Supports 128 or 256 bits
        blocksize:
            Size of crypted chunk. Supports 512, 1024, 4*1024, 16*1024, 64*1024 bits
        cipher:
            Chipher algoritm. Supports AES, Blowfish
        mode:
            Chipher mode. Supports CFB, PGP
        """
        ciphermode = str(ciphermode).upper()
        cipher = str(cipher).upper()
        assert MODE.has_key(ciphermode), 'Unknown cipher mode. Use CFB or PGP'
        assert CIPHER.has_key(cipher), 'Unknown cipher'
        assert blocksize in CHUNKSIZE, 'Bad blocksize'
        self.blocksize = int(blocksize)
        self.mode = MODE.get(ciphermode, None)
        self.cipher = CIPHER.get(cipher)[0]
        self.ivsize = CIPHER.get(cipher)[1]

        if not key:
            # If the password is callable we'll assume it will return the
            # real password.
            try:
                password = password()
            except TypeError:
                pass
            # the crypto algorithms are unicode unfriendly
            if isinstance(password, unicode):
                password = password.encode('utf8')
            if len(str(password)) == 0:
                print >>sys.stderr, 'Password is empty!'
            self.key = pass2key(password, message=login, keysize=keysize)
        else:
            assert len(str(key)) % int(keysize), 'Key is not %s long' % keysize
            self.key = key

    

    def GetIV(self):
        'Prepare IV'
        return os.urandom(self.ivsize)

    def EncChunk(self, chunk, iv=''):
        #assert hasattr(self.cipher, 'new') or type(self.cipher) == 'module', 'Bad cipher algorithm'
        encryptor = self.cipher.new(self.key,
                                    self.mode,
                                    iv)
        return encryptor.encrypt(chunk)

    def DecChunk(self, chunk, iv=''):
        #assert hasattr(self.cipher, 'new') or type(self.cipher) == 'module', 'Bad cipher algorithm'
        decryptor = self.cipher.new(self.key,
                                    self.mode,
                                    iv)
        return decryptor.decrypt(chunk)

    def GetMeta(self, infile):
        assert isinstance(infile, int), 'Bad file descryptor for reading meta'
        os.lseek(infile, 0, 0)
        iv = os.read(infile, self.ivsize)
        return iv

    def WriteMeta(self, infile, iv):
        assert isinstance(infile, int), 'Bad file descryptor for writing meta'
        assert len(iv) == self.ivsize, 'Bad IV size'
        os.lseek(infile, 0, 0)
        ret = os.write(infile, iv)
        os.close(os.dup(infile))
        return ret # вернем количество записанных байт

    def DecryptFile(self, infile, outfile):
        '''Decrypts a file using usally AES (CFB mode) with the
        given key.
        infile - input encrypted file.
        outfile - output decrypted file. 
        '''
        assert os.access(infile, os.R_OK), 'Bad encrpted input file'
        assert isinstance(outfile, str), 'outfile must be a string'

        with open(infile, 'rb') as ifd:
            with open(outfile, 'wb') as ofd:
                ifdno = ifd.fileno()
                ofdno = ofd.fileno()
                iv = self.GetMeta(ifdno)
                # пропустим служебную информацию
                os.lseek(ifdno, self.ivsize, 0)
                while True:
                    chunk = os.read(ifdno, self.blocksize)
                    if not chunk:
                        break
                    os.write(ofdno,
                             self.DecChunk(chunk, iv=iv))
                size = os.lseek(ofdno, 0, 2)
        return size # вернем количество записанных байт

    def EncryptFile(self, infile, outfile):
        ''' Encrypts a file using usally AES (CFB mode) with the
        given key.
        infile - input decrypted file.
        outfile - output encrypted file. 
        '''
        assert isinstance(outfile, str), 'outfile must be a string'
        assert os.access(infile, os.R_OK), 'Bad clear input file'

        iv = self.GetIV()

        with open(infile, 'rb') as ifd:
            with open(outfile, 'wb') as ofd:
                ifdno = ifd.fileno()
                ofdno = ofd.fileno()
                self.WriteMeta(ofdno, iv)
                # пропустим служебную информацию
                os.lseek(ofdno, self.ivsize, 0)
                while True:
                    chunk = os.read(ifdno, self.blocksize)
                    if not chunk:
                        break
                    os.write(ofdno, self.EncChunk(chunk, iv=iv))
                size = os.lseek(ofdno, 0, 2)
        return size # вернем количество записанных байт


### TESTS<<<
def tests():
    import unittest, tempfile
    class CryptoTests(unittest.TestCase):
        def testPass2Key(self):
            pw = ('123456', 1, os.urandom(10200))
            l = ('qwerty', 1, os.urandom(102400))
            for ks in KEYSIZE.keys():
                for login in l:
                    for password in pw:
                        self.assertIsInstance(pass2key(password, message=login, keysize=ks),
                                              basestring)

        def testCryptsAES(self):
            password = '123456'
            login = 'qwerty'
            cipher = 'AES'
            for mode in MODE.keys():
                for blocksize in CHUNKSIZE:
                    for keysize in KEYSIZE.keys():
                        m = Crypts(password,
                                    login=login,
                                    blocksize=blocksize,
                                    keysize=keysize,
                                    cipher=cipher,
                                    ciphermode=mode)
                        iv = m.GetIV()
                        string = ('',
                                os.urandom(20),
                                os.urandom(m.blocksize),
                                os.urandom(m.blocksize + 20),
                                os.urandom(m.blocksize * 7))
                        for s in string:
                            eraw = m.EncChunk(s, iv=iv)
                            self.assertEqual(m.DecChunk(eraw, iv=iv), s)

        def testCryptsDefaults(self):
            password = '123456'
            login = 'qwerty'
            m = Crypts(password=password, login=login)
            iv = m.GetIV()
            string = ('',
                    os.urandom(20),
                    os.urandom(m.blocksize),
                    os.urandom(m.blocksize + 20),
                    os.urandom(m.blocksize * 7))
            for s in string:
                eraw = m.EncChunk(s, iv=iv)
                self.assertEqual(m.DecChunk(eraw, iv=iv), s)

        def testCryptsBlowfish(self):
            password = '123456'
            login = 'qwerty'
            cipher = 'Blowfish'
            for mode in MODE:
                for blocksize in CHUNKSIZE:
                    for keysize in KEYSIZE.keys():
                        m = Crypts(password,
                                login=login,
                                blocksize=blocksize,
                                keysize=keysize,
                                cipher=cipher,
                                ciphermode=mode)
                        iv = m.GetIV()
                        string = ('',
                                os.urandom(20),
                                os.urandom(m.blocksize),
                                os.urandom(m.blocksize + 20),
                                os.urandom(m.blocksize * 7))
                        for s in string:
                            eraw = m.EncChunk(s, iv=iv)
                            self.assertEqual(m.DecChunk(eraw, iv=iv), s)

        def testCryptsFile(self):
            origfile = tempfile.mktemp(prefix='pyorig')
            encfile = tempfile.mktemp(prefix='pyenc')
            decfile = tempfile.mktemp(prefix='pydec')
            password = '123456'
            login = 'qwerty'
            for cipher in CIPHER.keys():
                for mode in MODE.keys():
                    for blocksize in CHUNKSIZE:
                        for keysize in KEYSIZE.keys():
                            m = Crypts(password,
                                        login=login,
                                        blocksize=blocksize,
                                        keysize=keysize,
                                        cipher=cipher,
                                        ciphermode=mode)
                            string = ('',
                                    os.urandom(20),
                                    os.urandom(m.blocksize),
                                    os.urandom(m.blocksize + 20),
                                    os.urandom(m.blocksize * 7))
                            for s in string:
                                with open(origfile, 'wb') as ifile:
                                    ifile.write(s)
                                m.EncryptFile(origfile, encfile)
                                m.DecryptFile(encfile, decfile)
                                with open(origfile, 'rb') as orig:
                                    with open(decfile, 'rb') as ofile:
                                        self.assertEqual(ofile.read(), orig.read())
            os.unlink(origfile)
            os.unlink(encfile)
            os.unlink(decfile)

        def testBigFile(self):
            password = '123456'
            login = 'qwerty'
            keysize = '128'
            blocksize = '4096'
            mode = 'CFB'
            origfile = tempfile.mktemp(prefix='pyorig')
            encfile = tempfile.mktemp(prefix='pyenc')
            decfile = tempfile.mktemp(prefix='pydec')

            
            for cipher in CIPHER.keys():
                m = Crypts(password,
                            login=login,
                            blocksize=blocksize,
                            keysize=keysize,
                            cipher=cipher,
                            ciphermode=mode)
                with open('/dev/zero', 'rb') as ifile:
                    with open(origfile, 'wb') as ofile:
                        ofile.write(ifile.read(100 * 2**20))

                m.EncryptFile(origfile, encfile)
                m.DecryptFile(encfile, decfile)
                with open(origfile, 'rb') as orig:
                    with open(decfile, 'rb') as ofile:
                        self.assertEqual(ofile.read(), orig.read())

            os.unlink(origfile)
            os.unlink(encfile)
            os.unlink(decfile)


    suite = unittest.TestSuite()
    suite.addTest(unittest.makeSuite(CryptoTests))
    unittest.TextTestRunner(verbosity=2).run(suite)
### TESTS>>>


def test1():
    import time
    password = '123456'
    login = 'qwerty'
    file = '/tmp/zero.100m'
    with open('/dev/zero', 'rb') as zero:
        with open(file, 'wb') as f:
            f.write(zero.read(100 * 2 ** 20))  # write 100Mb of zeros
    t1 = time.time()
    m = Crypts(password,
            login=login,
            blocksize='4096')
    t2 = time.time()
    t_init = t2 - t1
    m.EncryptFile(file, file + '.enc')
    t3 = time.time()
    t_enc = t3 - t2
    m.DecryptFile(file + '.enc', file + '.dec')
    t4 = time.time()
    t_dec = t4 - t3
    
    print 'Times:'
    print 't_init:', t_init
    print 't_enc:', t_enc
    print 't_dec:', t_dec

if __name__ == '__main__':
    test1()
